#!/usr/bin/env python
# $Revision: 2.1 $

import string, os, sys, stat, re

# Create a TRACE file if none exists
def init_trace():
    try:
        os.stat("TRACE")
    except OSError:
        os.system("rm -f TRACE")
        os.mkfifo("TRACE")

# Remove TRACE if it exists
def end_trace():
    if stat.S_ISFIFO(os.stat("TRACE")[stat.ST_MODE]):
        os.remove("TRACE")
    sys.exit(0)

def hex_conv(s):
    return int(s, 16)

# Load symbols
def load_symbols(prog):
    pattern = re.compile("[^:]+:\s*([0-9a-fA-F]+)? (\w) (\w+)")
    f = os.popen("nm -o "+prog)
    sym = {}
    for line in f.readlines():
        match = pattern.match(line)
        if not match or match.groups()[0] == None:
            continue
        addr, kind, name = match.groups()
        addr = hex_conv(addr)
        if name[0] == '_' and not name.startswith("__cyg_profile_func_"):
            name = name[1:]
        sym[addr]=name

    if "__cyg_profile_func_enter" not in sym.values():
        print("Error: "+prog+" doesn't appear to be instrumented. ")
        end_trace()

    f.close()
    return sym


# Print a program's trace
def trace(prog):
    sym = load_symbols(prog)
    trace_file = os.path.abspath("TRACE")
    tr = open(trace_file, 'r')

    printer = Printer()

    while True:
        line = tr.readline()
        if line:
            s = line.split()

            if s[0]=="EXIT":
                end_trace()

            s1 = hex_conv(s[1])

            if s1 in sym.keys():
                name = sym[s1]
            else:
                name = "??"

            printer.printFunction(s[0], name)
    tr.close()

class Printer:

    def __init__(self):
        self.level = 0
        self.prevname = ""
        self.count = 0

    def incr_count(self):
        self.count = self.count + 1

    def incr_level(self):
        self.level = self.level + 1

    def decr_level(self):
        self.level = self.level - 1

    def print_func(self, name):
        print ("\n   " + self.level*"|\t" + "\--" + name, )

    def print_totals(self):
        print ("(total: %d times)" % (self.count+1), )

    def do_func(self, name):
        if self.prevname == name:
            # was the counter counting?
            if self.count == 0:
                self.print_func(name)
            self.incr_count()
        else:
            # New name received. Was the counter counting?
            if self.count > 0:
                self.print_totals()
                self.count=0

            self.print_func(name)

    def printFunction(self, status, name):
        if status == "enter":
            self.do_func(name)
            # Check if new name compared to previous input
            self.incr_level()
            self.prevname = name
        else:
            self.decr_level()
            if name == "main":
                if self.count > 0:
                    self.print_totals()
                end_trace()

# Main
if __name__=="__main__":
    if len(sys.argv)<2:
        print("Usage: analysis <binary name>")
        sys.exit(1)

    init_trace()
    trace(sys.argv[1])
